#include "xtcp.h"
#include "gdbinfo.h"
#include "xip.h"

#include <stdlib.h>

#define LOCAL_WIN_SIZE 1024
#define XTCP_DATA_MAX_SIZE (XNET_CFG_PACKET_MAX_SIZE - sizeof(xether_hdr_t) - sizeof(xip_hdr_t) - sizeof(xtcp_hdr_t))

static xtcp_t tcp_socket[XTCP_CFG_MAX_TCP]; // TCP连接块

static uint32_t tcp_get_init_seq()
{
    // 随机生成一个seq
    return ((rand() << 16) + rand());
}

static xtcp_t *tcp_alloc()
{
    // 分配一个空闲的tcp块
    for (xtcp_t *tcp = tcp_socket, *end = &tcp_socket[XTCP_CFG_MAX_TCP]; tcp < end; tcp++)
    {
        if (tcp->state == XTCP_STATE_FREE)
        {
            tcp->state = XTCP_STATE_CLOSED;
            tcp->local_port = 0;
            tcp->remote_port = 0;
            tcp->remote_ip.addr = 0;
            tcp->handler = (xtcp_handler_t)0;
            tcp->remote_win = XTCP_MSS_DEFAULT;
            tcp->remote_mss = XTCP_MSS_DEFAULT;
            tcp->next_seq = tcp->unacked_seq = tcp_get_init_seq();
            tcp->next_ack = 0;
            return tcp;
        }
    }
    return (xtcp_t *)0;
}

static void tcp_free(xtcp_t *tcp)
{
    if (TCP_GDB_INFO)
    {
        char ip_str[32];
        ip_to_str(tcp->remote_ip.array, ip_str, XNET_IPV4_ADDR_SIZE);
        print_info("free tcp control block remote: %s port %d  local %d", ip_str, tcp->remote_port, tcp->local_port);
    }
    tcp->state = XTCP_STATE_FREE;
}

xnet_err_t tcp_send(xtcp_t *tcp, uint8_t flags)
{
    // 将缓冲区的数据发送出去
    // 发送tcp的包
    xnet_packet_t *packet;
    xtcp_hdr_t *tcp_hdr;
    xnet_err_t err;
    uint16_t data_size = tcp_buffer_wait_send_size(&tcp->tx_buf);
    uint16_t opt_size = (flags & XTCP_FLAG_SYN) ? 4 : 0; // 有syn，意味着发送tcp三次握手的第二次

    // 判断当前允许发送的字节量
    if (tcp->remote_win)
    {
        data_size = min(data_size, tcp->remote_win);
        data_size = min(data_size, tcp->remote_mss);
        if (data_size + opt_size > XTCP_DATA_MAX_SIZE)
        {
            data_size = XTCP_DATA_MAX_SIZE - opt_size;
        }
    }
    else
    {
        data_size = 0; // 窗口为0，不允许发数据，但允许FIN+SYN等
    }

    packet = xnet_alloc_for_send(data_size + opt_size + sizeof(xtcp_hdr_t));
    tcp_hdr = (xtcp_hdr_t *)packet->data;
    tcp_hdr->src_port = swap_order16(tcp->local_port);
    tcp_hdr->dest_port = swap_order16(tcp->remote_port);
    tcp_hdr->seq = swap_order32(tcp->next_seq);
    tcp_hdr->ack = swap_order32(tcp->next_ack);

    tcp_hdr->hdr_flags.all = 0;
    tcp_hdr->hdr_flags.hdr_len = (opt_size + sizeof(xtcp_hdr_t)) / 4;
    tcp_hdr->hdr_flags.flag = flags;
    tcp_hdr->hdr_flags.all = swap_order16(tcp_hdr->hdr_flags.all);

    tcp_hdr->window = swap_order16(LOCAL_WIN_SIZE);
    tcp_hdr->checksum = 0;
    tcp_hdr->urgent_ptr = 0;

    if (flags & XTCP_FLAG_SYN)
    {
        // 填充可选字段的mss
        uint8_t *opt_data = packet->data + sizeof(xtcp_hdr_t);
        opt_data[0] = XTCP_KIND_MSS;
        opt_data[1] = 4;
        *(uint16_t *)(opt_data + 2) = swap_order16(XTCP_MSS_DEFAULT);
    }

    tcp_buffer_read(&tcp->tx_buf, packet->data + opt_size + sizeof(xtcp_hdr_t), data_size);

    // 校验和
    tcp_hdr->checksum = checksum_peso(&my_ip, &tcp->remote_ip, XNET_PROTOCOL_TCP, (uint16_t *)packet->data, packet->size);
    tcp_hdr->checksum = tcp_hdr->checksum ? tcp_hdr->checksum : 0xffff;

    err = xip_out(XNET_PROTOCOL_TCP, &tcp->remote_ip, packet);
    if (err < 0)
    {
        return err;
    }

    if (tcp_hdr->hdr_flags.flag & (XTCP_FLAG_FIN | XTCP_FLAG_SYN))
    {
        tcp->next_seq++; // 占一个字节
    }
    tcp->next_seq += data_size;
    return XNET_ERR_OK;
}

xtcp_t *xtcp_open(xtcp_handler_t handler)
{
    xtcp_t *tcp = tcp_alloc();
    if (!tcp)
        return (xtcp_t *)0;
    tcp->state = XTCP_STATE_CLOSED;
    tcp->handler = handler;
    return tcp;
};

xnet_err_t xtcp_bind(xtcp_t *tcp, uint16_t local_port)
{
    // 将tcp和本地端口绑定
    xtcp_t *curr, *end;
    for (curr = tcp_socket, end = &tcp_socket[XTCP_CFG_MAX_TCP]; curr < end; curr++)
    {
        if ((curr != tcp) && (curr->local_port == local_port))
        {
            return XNET_ERR_BINDED;
        }
    }

    tcp->local_port = local_port;
    return XNET_ERR_OK;
};

xnet_err_t xtcp_listen(xtcp_t *tcp)
{
    // tcp监听
    tcp->state = XTCP_STATE_LISTEN;
    return XNET_ERR_OK;
};

xnet_err_t xtcp_close(xtcp_t *tcp)
{
    // 关闭tcp
    xnet_err_t err;

    if (tcp->state == XTCP_STATE_ESTABLISHED)
    {
        err = tcp_send(tcp, XTCP_FLAG_FIN | XTCP_FLAG_ACK);
        if (err < 0)
            return err;

        tcp->state = XTCP_STATE_FIN_WAIT_1;
    }
    else
    {
        tcp_free(tcp);
    }
    return XNET_ERR_OK;
};

static xtcp_t *tcp_find(xipaddr_t *remote_ip, uint16_t remote_port, uint16_t local_port)
{
    // 找到remote ip  remote port local port对应的控制块
    xtcp_t *tcp, *end;
    xtcp_t *res = (xtcp_t *)0;
    for (tcp = tcp_socket, end = &tcp_socket[XTCP_CFG_MAX_TCP]; tcp < end; tcp++)
    {
        if ((tcp->state == XTCP_STATE_FREE) || (tcp->local_port != local_port))
        {
            continue;
        }

        if (xipaddr_is_equal(*remote_ip, tcp->remote_ip.addr) && (remote_port == tcp->remote_port))
        {
            return tcp;
        }

        if (tcp->state == XTCP_STATE_LISTEN)
        {
            res = tcp; // 刚开始的时候，只创建了一个监听tcp的控制块，这里返回之后，会调用tcp_process_accept，初始化一个remote_ip为远程的tcp控制块
        }
    }
    return res;
}

static xnet_err_t tcp_send_reset(uint32_t remote_ack, uint16_t local_port, xipaddr_t *remote_ip, uint16_t remote_port)
{
    // 发送tcp复位包
    xnet_packet_t *packet = xnet_alloc_for_send(sizeof(xtcp_hdr_t));
    xtcp_hdr_t *tcp_hdr = (xtcp_hdr_t *)packet->data;

    tcp_hdr->src_port = swap_order16(local_port);
    tcp_hdr->dest_port = swap_order16(remote_port);
    tcp_hdr->seq = 0;                        // 固定为0即可
    tcp_hdr->ack = swap_order32(remote_ack); // 响应指定的发送ack，即对上次发送的包的回应
    tcp_hdr->hdr_flags.all = 0;
    tcp_hdr->hdr_flags.hdr_len = sizeof(xtcp_hdr_t) / 4;
    tcp_hdr->hdr_flags.flag = XTCP_FLAG_RST | XTCP_FLAG_ACK;
    tcp_hdr->hdr_flags.all = swap_order16(tcp_hdr->hdr_flags.all);
    tcp_hdr->window = 0;
    tcp_hdr->urgent_ptr = 0;
    tcp_hdr->checksum = 0;
    tcp_hdr->checksum = checksum_peso(&my_ip, remote_ip, XNET_PROTOCOL_TCP, (uint16_t *)packet->data, packet->size);
    tcp_hdr->checksum = tcp_hdr->checksum ? tcp_hdr->checksum : 0xFFFF;

    if (TCP_GDB_INFO)
    {
        char ip_str[32];
        ip_to_str(remote_ip, ip_str, XNET_IPV4_ADDR_SIZE);
        print_info("try to send reset tcp packet to %s port %d", ip_str, remote_port);
    }

    return xip_out(XNET_PROTOCOL_TCP, remote_ip, packet);
}

static void tcp_read_mss(xtcp_t *tcp, xtcp_hdr_t *tcp_hdr)
{
    // 从包里面读取mss ,并赋值给tcp里的remote_mss
    uint16_t opt_len = tcp_hdr->hdr_flags.hdr_len * 4 - sizeof(xtcp_hdr_t);

    if (opt_len == 0)
    {
        tcp->remote_mss = XTCP_MSS_DEFAULT;
    }
    else
    {
        uint8_t *opt_data = (uint8_t *)tcp_hdr + sizeof(xtcp_hdr_t);
        uint8_t *opt_end = opt_data + opt_len;

        while ((*opt_data != XTCP_KIND_END) && (opt_data < opt_end))
        {
            if ((*opt_data++ == XTCP_KIND_MSS) && (*opt_data++ == 4))
            {
                tcp->remote_mss = swap_order16(*(uint16_t *)opt_data);
                return;
            }
        }
    }
}

static void tcp_process_accept(xtcp_t *listen_tcp, xipaddr_t *remote_ip, xtcp_hdr_t *tcp_hdr)
{

    // 为连接分配tcp控制块，这个时候，当packet 有 syn标志时，标志着对方已经发送了三次握手的第一次
    uint16_t hdr_flag = tcp_hdr->hdr_flags.all;

    if (hdr_flag & XTCP_FLAG_SYN)
    {
        xtcp_t *new_tcp = tcp_alloc();
        if (!new_tcp)
        {
            // 超过了最大连接数量
            return;
        }
        new_tcp->state = XTCP_STATE_SYNC_RECVD;
        new_tcp->local_port = listen_tcp->local_port;
        new_tcp->remote_port = tcp_hdr->src_port;
        new_tcp->handler = listen_tcp->handler;
        new_tcp->next_ack = tcp_hdr->seq + 1; // syn占一个
        new_tcp->remote_ip.addr = remote_ip->addr;
        new_tcp->next_seq = tcp_get_init_seq(); // 期待下次一发送的序列号
        new_tcp->remote_win = tcp_hdr->window;
        tcp_buffer_init(&new_tcp->tx_buf);
        tcp_buffer_init(&new_tcp->rx_buf);
        tcp_read_mss(new_tcp, tcp_hdr); // 读选项，主要是mss
        xnet_err_t err = tcp_send(new_tcp, XTCP_FLAG_SYN | XTCP_FLAG_ACK);

        if (err < 0)
        {
            tcp_free(new_tcp);
            return;
        }
        else
        {
            if (TCP_GDB_INFO)
            {
                char ip_str[32];
                ip_to_str(remote_ip, ip_str, XNET_IPV4_ADDR_SIZE);
                print_info("init tcp control block remote: %s port %d  local %d", ip_str, tcp_hdr->src_port, listen_tcp->local_port);
            }
        }
    }
    else
    {
        if (TCP_GDB_INFO)
        {
            print_info("no connection between remote port %d and local port %d", tcp_hdr->dest_port, listen_tcp->local_port);
        }
        // 发送复位信号
        tcp_send_reset(tcp_hdr->seq, listen_tcp->local_port, remote_ip, tcp_hdr->src_port);
    }
}

xnet_err_t tcp_close(xtcp_t *tcp)
{
    // 主动关闭tcp连接
    if (tcp->state == XTCP_STATE_ESTABLISHED)
    {
        xnet_err_t err = tcp_send(tcp, XTCP_FLAG_FIN | XTCP_FLAG_ACK);
        if (err < 0)
            return err;
        tcp->state = XTCP_STATE_FIN_WAIT_1;

        if (TCP_GDB_INFO)
        {
            print_info("release the connection actively by server");
        }
    }
    else
    {
        // 关闭不处于establish的tcp
        tcp_free(tcp);
    }
    return XNET_ERR_OK;
};

void xtcp_init()
{
    memset(tcp_socket, 0, sizeof(tcp_socket));
}

void xtcp_in(xipaddr_t *remote_ip, xnet_packet_t *packet)
{
    xtcp_hdr_t *tcp_hdr = (xtcp_hdr_t *)packet->data;

    if (packet->data < sizeof(xtcp_hdr_t))
    {
        return;
    }

    uint16_t pre_checksum = tcp_hdr->checksum;
    tcp_hdr->checksum = 0;
    if (pre_checksum != 0)
    {
        uint16_t checksum = checksum_peso(remote_ip, &my_ip, XNET_PROTOCOL_TCP, (uint16_t *)packet->data, packet->size);
        if (checksum == 0)
            checksum = 0xffff;
        if (checksum != pre_checksum)
        {

            if (NEED_UDP_GDB)
            {
                print_info("tcp receive a wrong checksum packet");
            }
            return;
        }
    }

    // 解析包里的参数，省着一遍遍的swap
    tcp_hdr->src_port = swap_order16(tcp_hdr->src_port);
    tcp_hdr->dest_port = swap_order16(tcp_hdr->dest_port);
    tcp_hdr->hdr_flags.all = swap_order16(tcp_hdr->hdr_flags.all);
    tcp_hdr->seq = swap_order32(tcp_hdr->seq);
    tcp_hdr->ack = swap_order32(tcp_hdr->ack);
    tcp_hdr->window = swap_order16(tcp_hdr->window);

    // 找到对应的tcp
    xtcp_t *tcp = tcp_find(remote_ip, tcp_hdr->src_port, tcp_hdr->dest_port);

    if (tcp == (xtcp_t *)0)
    {
        if (TCP_GDB_INFO)
        {
            print_info("no listening program");
        }
        // 未找到对应的tcp控制块，发送复位包
        tcp_send_reset(tcp_hdr->seq + 1, tcp_hdr->dest_port, remote_ip, tcp_hdr->src_port);
        return;
    }

    tcp->remote_win = tcp_hdr->window;
    if (tcp->state == XTCP_STATE_LISTEN)
    {
        tcp_process_accept(tcp, remote_ip, tcp_hdr);
        return;
    }

    // 序号不一致，可能要进行重发
    if (tcp_hdr->seq != tcp->next_ack)
    {
        if (TCP_GDB_INFO)
        {
            print_info("wrong seq packet");
        }
        tcp_send_reset(tcp_hdr->seq + 1, tcp_hdr->dest_port, remote_ip, tcp_hdr->src_port);
        return;
    }

    remove_header(packet, tcp_hdr->hdr_flags.hdr_len * 4);
    switch (tcp->state)
    {
    case XTCP_STATE_SYNC_RECVD:
        if (tcp_hdr->hdr_flags.flag & XTCP_FLAG_ACK)
        {
            tcp->next_seq = tcp->unacked_seq = tcp_hdr->ack;
            tcp->state = XTCP_STATE_ESTABLISHED;
            tcp->handler(tcp, XTCP_CONN_CONNECTED);
        }
        break;
    case XTCP_STATE_ESTABLISHED:
        if (tcp_hdr->hdr_flags.flag & (XTCP_FLAG_ACK | XTCP_FLAG_FIN))
        {

            if (tcp_hdr->hdr_flags.flag & XTCP_FLAG_ACK)
            {
                if ((tcp->unacked_seq <= tcp_hdr->ack) && (tcp_hdr->ack <= tcp->next_seq))
                {
                    uint16_t curr_ack_size = tcp_hdr->ack - tcp->unacked_seq;
                    tcp_buffer_add_acked_count(&tcp->tx_buf, curr_ack_size);
                    tcp->unacked_seq += curr_ack_size;
                }
                else if (TCP_GDB_INFO)
                {
                    print_info("ack out of range packet ack %d  unacked_seq = %d next_seq=%d", tcp_hdr->ack, tcp->unacked_seq, tcp->next_seq);
                }
            }

            // packet.data --> tcp->rx_buffer
            uint16_t read_size = tcp_buffer_write(&tcp->rx_buf, packet->data, packet->size);
            tcp->next_ack += read_size;

            if (tcp_hdr->hdr_flags.flag & XTCP_FLAG_FIN)
            {
                // 收到了对方请求关闭的包
                tcp->state = XTCP_STATE_LAST_ACK; // 等待对方发送的ack
                tcp->next_seq = tcp_hdr->ack;
                tcp->next_ack = tcp_hdr->seq + 1; // fin占一个
                tcp_send(tcp, XTCP_FLAG_FIN | XTCP_FLAG_ACK);
            }
            else if (read_size)
            {
                tcp_send(tcp, XTCP_FLAG_ACK);
                // 有数据，把tcp给对应的回调函数
                tcp->handler(tcp, XTCP_CONN_DATA_RECV);
            }
            else if (tcp_buffer_wait_send_size(&tcp->tx_buf))
            {
                tcp_send(tcp, XTCP_FLAG_ACK);
            }
        }

        break;
    case XTCP_STATE_LAST_ACK:
        if (tcp_hdr->hdr_flags.flag & XTCP_FLAG_ACK)
        {
            tcp->handler(tcp, XTCP_CONN_CLOSED);
            tcp_free(tcp);
        }
    case XTCP_STATE_FIN_WAIT_1:
        //  主动关闭后的第一个状态
        if ((tcp_hdr->hdr_flags.flag & (XTCP_FLAG_FIN | XTCP_FLAG_ACK)) == (XTCP_FLAG_FIN | XTCP_FLAG_ACK))
        {
            // 收到fin和ack，直接关闭
            tcp_free(tcp);
        }
        else if (tcp_hdr->hdr_flags.flag & XTCP_FLAG_ACK)
        {
            // 收到fin，等待对方fin
            tcp->state = XTCP_STATE_FIN_WAIT_2;
        }
        break;
    case XTCP_STATE_FIN_WAIT_2:
        if (tcp_hdr->hdr_flags.flag & XTCP_FLAG_FIN)
        {
            tcp->next_ack++;
            tcp_send(tcp, XTCP_FLAG_ACK);
            tcp_free(tcp);
        }
    default:
        break;
    }
};